import os
import numpy as np
from shapely.geometry import Polygon
from terminaltables import AsciiTable    
import pandas as pd

try:
    from cvio import cvio
except:
    from .cvio import cvio

def print_summary(table):
    table = AsciiTable(table)
    table.inner_footing_row_border = True
    print(table.table)

def load_images_shapes(srcs):
    image2instances = {}
    for src in srcs:
        ann = cvio.load_ann(src)
        image = os.path.basename(src)
        shapes = ann['shapes']
        image2instances[image] = shapes
    return image2instances

def polygon_iou_matrix(gtanns, ptanns):
    ioum = []
    for gtann in gtanns:
        points = gtann['points']
        if gtann['shape_type'] == 'rectangle':
            assert len(points) == 2
            (x1, y1), (x2, y2) = points
            points = [[x1, y1], [x1, y2], [x2, y2], [x2, y1]]
        if len(points) <= 3:
            # print('Exception', ptann)
            continue        
        gtply = Polygon(points).buffer(0)
        ioum.append([])
        for ptann in ptanns:
            points = ptann['points']
            if len(points) <= 3:
                # print('Exception', ptann)
                continue
            ptply = Polygon(points).buffer(0)
            inter = ptply.intersection(gtply).area
            union = ptply.union(gtply).area
            iou = inter / union
            ioum[-1].append(iou)
    return np.array(ioum)

def eval_polygon_map(gtsrc, ptsrc, iouthr=0.5, xls=None):
    gtsrcs = cvio.load_ext_list(gtsrc)
    ptsrcs = cvio.load_ext_list(ptsrc)

    gtanns = load_images_shapes(gtsrcs)
    ptanns = load_images_shapes(ptsrcs)

    map = {}
    N = len(gtanns)
    for i, (name, gtann) in enumerate(gtanns.items()):
        print('[%d/%d]' % (i + 1, N), name)
        ptann = ptanns[name] if name in ptanns else []
        ioum = polygon_iou_matrix(gtann, ptann)
        dets = []
        for j, iou in enumerate(ioum):
            gtlabel = gtann[j]['label']
            if gtlabel not in map:
                map[gtlabel] = dict(tp=0, fp=0, fn=0, gts=0, dets=0)
            map[gtlabel]['gts'] += 1
            if len(iou) == 0:
                map[gtlabel]['fn'] += 1
                continue
            idx = np.argmax(iou)
            iou = iou[idx]
            if iou < iouthr:
                map[gtlabel]['fn'] += 1
                continue
            ptlabel = ptann[idx]['label']
            if gtlabel != ptlabel:
                map[gtlabel]['fp'] += 1
                continue
            map[gtlabel]['tp'] += 1
            map[gtlabel]['dets'] += 1
            dets.append(idx)
        for i, iou in enumerate(ioum.T):
            if i in (dets):
                continue
            ptlabel = ptann[i]['label']
            if ptlabel in map:
                map[ptlabel]['fp'] += 1
                map[ptlabel]['dets'] += 1

    map, mapdf, table = parse_map(map)
    print_summary(table)
    if xls in (None, ''):
        xls = os.path.join(ptsrc, 'mask_eval.xlsx')
    save_map_to_excel(mapdf, xls)

def save_map_to_excel(mapdf, xls):
    mapdf = pd.DataFrame(mapdf)
    print('评估结果保存至', xls)
    mapdf.to_excel(xls, sheet_name='mAP', index=None)

def parse_map(map):
    
    mapdf = dict(classes=[], gts=[], dets=[], tp=[], fp=[], fn=[], recall=[], precision=[])
    def update_mapdf(label, gts, dets, tp, fp, fn, recall, precision):
        mapdf['classes'] += [label]
        mapdf['gts'] += [gts]
        mapdf['dets'] += [dets]
        mapdf['tp'] += [tp]
        mapdf['fp'] += [fp]
        mapdf['fn'] += [fn]
        mapdf['recall'] += [recall]
        mapdf['precision'] += [precision]

    gts = 0
    dets = 0
    tp, fp, fn = 0, 0, 0
    recall = []
    precision = []
    headers = ['classes', 'gts', 'dets', 'tp', 'fp', 'fn', 'recall', 'precision']
    table = [headers]    
    for label, item in map.items():
        tp += item['tp']
        fp += item['fp']
        fn += item['fn']
        gts += item['gts']
        dets += item['dets']
        item['recall'] = item['tp'] / (item['tp'] + item['fn'])
        item['precision'] = item['tp'] / (item['tp'] + item['fp'])
        recall += [item['recall']]
        precision += [item['precision']]
        update_mapdf(label, item['gts'], item['dets'], item['tp'], item['fp'], item['fn'], item['recall'], item['precision'])
        row = [label, item['gts'], item['dets'], item['tp'], item['fp'], item['fn'], item['recall'], item['precision']]
        table.append(row)
    update_mapdf('mAP', gts, dets, tp, fp, fn, np.mean(recall), np.mean(precision))
    table.append(['mAP', gts, dets, tp, fp, fn, np.mean(recall), np.mean(precision)])
    return map, mapdf, table


if __name__ == '__main__':
    ptsrc = r'G:\data\datasets\drink\daiding_101\base_test\yy\queryinst_swin_tiny_mask_bottle'
    gtsrc = r'G:\data\datasets\drink\daiding_101\base_test\yy'
    eval_polygon_map(gtsrc, ptsrc)